#!/usr/bin/env python3
import os
import json
from datetime import datetime
from typing import List, Dict, Any
from langchain.schema import Document
from database.mysql_manager import db_manager

class OrderKnowledgeBase:
    """
    订单知识库类，用于从MySQL数据库提取订单数据并准备用于向量存储
    """
    
    def __init__(self):
        # 统一使用订单知识库集合名称
        self.order_collection_name = "orders_knowledge_base"
    
    def fetch_orders_from_mysql(self, limit: int = None) -> List[Dict[str, Any]]:
        """从MySQL数据库获取订单数据"""
        try:
            if not db_manager.connect():
                print("无法连接到MySQL数据库")
                return []
            
            # 构建SQL查询
            query = """
                SELECT o.id, o.order_no, o.customer_name, o.customer_email, 
                       o.total_amount, o.status, o.created_at, o.updated_at,
                       GROUP_CONCAT(CONCAT('{"product_name":"', oi.product_name, '","quantity":', oi.quantity, ',"unit_price":', oi.unit_price, '}') SEPARATOR ',') as items
                FROM orders o
                LEFT JOIN order_items oi ON o.id = oi.order_id
                GROUP BY o.id
                ORDER BY o.created_at DESC
            """
            
            if limit:
                query += f" LIMIT {limit}"
            
            # 执行查询
            results = db_manager.execute_query(query)
            
            # 处理结果，将items字符串解析为JSON数组
            for result in results:
                if result['items']:
                    try:
                        result['items'] = json.loads(f"[{result['items']}]")
                    except json.JSONDecodeError:
                        result['items'] = []
                else:
                    result['items'] = []
            
            return results
            
        except Exception as e:
            print(f"从MySQL获取订单数据时出错: {e}")
            return []
        finally:
            db_manager.disconnect()
    
    def convert_orders_to_documents(self, orders: List[Dict[str, Any]]) -> List[Document]:
        """将订单数据转换为Langchain的Document对象"""
        documents = []
        
        for order in orders:
            # 构建订单的文本表示
            order_text = self._format_order_as_text(order)
            
            # 创建Document对象，包含订单文本和元数据
            document = Document(
                page_content=order_text,
                metadata={
                    "order_id": order["id"],
                    "order_no": order["order_no"],
                    "customer_name": order["customer_name"],
                    "status": order["status"],
                    "created_at": order["created_at"].isoformat() if isinstance(order["created_at"], datetime) else order["created_at"],
                    "total_amount": float(order["total_amount"]),
                    "source": "mysql_order_data"
                }
            )
            
            documents.append(document)
        
        return documents
    
    def _format_order_as_text(self, order: Dict[str, Any]) -> str:
        """将单个订单格式化为文本字符串"""
        # 格式化订单基本信息
        text_parts = [
            f"订单号: {order['order_no']}",
            f"客户: {order['customer_name']}",
            f"联系邮箱: {order['customer_email']}",
            f"订单状态: {order['status']}",
            f"总金额: {float(order['total_amount']):.2f}元",
            f"创建时间: {order['created_at']}"
        ]
        
        # 添加订单项信息
        text_parts.append("\n订单项:")
        for item in order.get('items', []):
            item_total = float(item['quantity']) * float(item['unit_price'])
            text_parts.append(f"- {item['product_name']}: {item['quantity']}件 × {float(item['unit_price']):.2f}元 = {item_total:.2f}元")
        
        return "\n".join(text_parts)
    
    def update_order_knowledge_base(self, vector_store_manager, limit: int = None):
        """更新订单知识库，将MySQL中的订单数据同步到向量存储"""
        try:
            print(f"开始更新订单知识库，集合名称: {self.order_collection_name}")
            
            # 从MySQL获取订单数据
            orders = self.fetch_orders_from_mysql(limit)
            print(f"从MySQL获取到 {len(orders)} 个订单")
            
            if not orders:
                print("没有获取到订单数据，跳过更新")
                return False
            
            # 将订单数据转换为Document对象
            documents = self.convert_orders_to_documents(orders)
            print(f"转换为 {len(documents)} 个Document对象")
            
            # 检查是否已存在订单知识库集合，如果存在则删除（为了完全更新）
            existing_collections = vector_store_manager.get_collections()
            if self.order_collection_name in existing_collections:
                print(f"删除已存在的订单知识库集合: {self.order_collection_name}")
                vector_store_manager.delete_collection(self.order_collection_name)
            
            # 将文档添加到向量存储
            # 注意这里调用方式可能需要根据vector_store_manager的实现调整
            try:
                # 尝试使用参数名方式调用
                chunks_added = vector_store_manager.add_documents(
                    documents=documents, 
                    collection_name=self.order_collection_name
                )
                # 处理返回结果
                if isinstance(chunks_added, dict) and chunks_added.get("status") == "success":
                    print(f"成功添加文档到订单知识库")
                    return True
                elif isinstance(chunks_added, int):
                    print(f"成功添加 {chunks_added} 个文档块到订单知识库")
                    return chunks_added > 0
                else:
                    print(f"添加文档到订单知识库失败")
                    return False
            except TypeError:
                # 如果第一种调用方式失败，尝试使用位置参数方式
                try:
                    result = vector_store_manager.add_documents(self.order_collection_name, documents)
                    if isinstance(result, dict) and result.get("status") == "success":
                        print(f"成功添加文档到订单知识库")
                        return True
                    else:
                        print(f"添加文档到订单知识库失败")
                        return False
                except Exception as inner_e:
                    print(f"添加文档到向量存储时出错: {inner_e}")
                    return False
            
        except Exception as e:
            print(f"更新订单知识库时出错: {e}")
            return False
    
    def search_order_knowledge(self, vector_store_manager, query: str, top_k: int = 3) -> List[Dict[str, Any]]:
        """
        在订单知识库中搜索相关信息
        """
        try:
            # 确保集合存在
            existing_collections = vector_store_manager.get_collections()
            if self.order_collection_name not in existing_collections:
                print(f"订单知识库集合 {self.order_collection_name} 不存在")
                return []
            
            # 执行向量搜索
            try:
                # 尝试使用参数名方式调用
                search_result = vector_store_manager.search(
                    query=query,
                    collection_name=self.order_collection_name,
                    top_k=top_k
                )
            except TypeError:
                # 如果第一种调用方式失败，尝试使用位置参数方式
                search_result = vector_store_manager.search(self.order_collection_name, query, top_k)
            
            # 处理搜索结果
            if isinstance(search_result, dict):
                if search_result.get("status") != "success":
                    print(f"搜索订单知识库失败：{search_result.get('message', '未知错误')}")
                    return []
                results = search_result.get("results", [])
            else:
                # 如果返回的直接是结果列表
                results = search_result
            
            # 格式化搜索结果
            formatted_results = []
            for item in results:
                # 处理不同格式的结果对象
                if hasattr(item, 'page_content'):
                    # 如果是Document对象
                    content = item.page_content
                    metadata = item.metadata
                    score = getattr(item, 'metadata', {}).get('score', 0)
                elif isinstance(item, dict):
                    # 如果是字典
                    content = item.get("content", item.get("page_content", ""))
                    metadata = item.get("metadata", {})
                    score = item.get("score", 0)
                else:
                    # 未知格式
                    content = str(item)
                    metadata = {}
                    score = 0
                
                formatted_result = {
                    "order_id": metadata.get("order_id"),
                    "customer_name": metadata.get("customer_name"),
                    "order_date": metadata.get("created_at", metadata.get("order_date")),
                    "total_amount": metadata.get("total_amount"),
                    "content": content,
                    "score": score
                }
                formatted_results.append(formatted_result)
            
            print(f"在订单知识库中找到 {len(formatted_results)} 条相关结果")
            return formatted_results
        
        except Exception as e:
            print(f"搜索订单知识库时出错：{str(e)}")
            return []
            
    def generate_order_answer(self, llm_wrapper, query: str, search_results: List[Dict[str, Any]]):
        """
        使用LLM基于搜索结果生成回答
        """
        try:
            # 构建上下文信息
            context = "\n".join([
                f"订单ID: {result.get('order_id')}, 客户: {result.get('customer_name')}, "
                f"日期: {result.get('order_date')}, 总金额: {result.get('total_amount')}, "
                f"详情: {result.get('content')}"
                for result in search_results
            ])
            
            # 构建系统提示，特别针对订单查询进行优化
            system_prompt = """
            你是一个订单管理智能助手，根据提供的订单上下文信息回答用户的问题。
            请确保你的回答完全基于提供的上下文信息，不要添加外部知识。
            如果上下文信息不足以回答问题，请明确表示无法回答。
            回答时请保持专业、准确，并尽量详细地提供订单相关信息。
            """
            
            # 构建用户提示
            prompt = f"""请根据提供的订单上下文信息，回答用户的问题。
如果无法从上下文获取足够信息，请说明无法回答。

上下文信息：
{context}

用户问题：
{query}

请用中文回答。"""
            
            # 使用LLM生成回答，尝试不同的参数调用方式
            try:
                answer_result = llm_wrapper.generate_answer(
                    prompt=prompt,
                    system_prompt=system_prompt
                )
            except TypeError:
                # 如果参数不匹配，尝试直接传递查询和上下文
                answer_result = llm_wrapper.generate_answer(query, search_results)
            
            return answer_result
        
        except Exception as e:
            print(f"生成订单回答时出错：{str(e)}")
            return {
                "answer": f"生成回答时出错：{str(e)}",
                "error": True
            }

# 创建全局实例
def get_order_knowledge_base():
    """获取订单知识库的单例实例"""
    if not hasattr(get_order_knowledge_base, '_instance'):
        get_order_knowledge_base._instance = OrderKnowledgeBase()
    return get_order_knowledge_base._instance

# 当直接运行此脚本时执行更新操作
if __name__ == "__main__":
    from backend.vector_store.vector_store import VectorStoreManager
    from backend.llm.llm_wrapper import LLMWrapper
    
    print("订单知识库更新工具")
    print("================")
    
    try:
        # 初始化向量存储管理器
        vector_store = VectorStoreManager()
        
        # 初始化订单知识库
        order_kb = OrderKnowledgeBase()
        
        # 更新订单知识库
        success = order_kb.update_order_knowledge_base(vector_store)
        
        if success:
            print("\n订单知识库更新成功!")
        else:
            print("\n订单知识库更新失败!")
            
    except Exception as e:
        print(f"执行过程中出错: {e}")